
from __future__ import annotations
from typing import  Generic, TypeVar, Callable, Optional, List, Set, Deque, Dict
from heapq import heappush, heappop

T = TypeVar('T')

# 栈，后进先出，用于DFS
class Stack(Generic[T]):
    def __init__(self) -> None:
        self._container: List[T] = []

    @property
    def empty(self) -> bool:
        return not self._container  # not is true for empty container

    def push(self, item: T) -> None:
        self._container.append(item)

    def pop(self) -> T:
        return self._container.pop() #LIFO

    def __repr__(self) -> str:
        return repr(self._container)

# 引入自身需要导入annotations
class Node(Generic[T]):
    def __init__(self, state: T, parent: Optional[Node], cost: float = 0.0, heuristic: float = 0.0) -> None:
        self.state: T = state
        self.parent: Optional[Node] = parent
        self.cost: float = cost
        self.heuristic: float = heuristic

    def __lt__(self, other: Node) -> bool:
        return (self.cost + self.heuristic) < (other.cost + other.heuristic)

# 搜索：深度优先，朝向一个方向，尽快找到目标
# 方向由successors函数返回的列表顺序决定（最后的节点优先弹出）
# 搜索的三个关键因素：起始点，下一步可能，是否完成
# 内部执行的关键数据：已经探索的部分，下一步需要探索的部分
# 方式的关键点：下一步的选择方法
def dfs(initial: T, goal_test: Callable[[T], bool], successors: Callable[[T], List[T]]) -> Optional[Node[T]]:
    # 代码解释：py的类型：Callable[[T], List[T]]： 这表示的是一个接受一个类型为 T 参数的可调用对象，它返回一个包含 T 类型元素的列表（List[T]）

    # frontier is where we've yet to go
    frontier: Stack[Node[T]] = Stack()
    # 先传入起点
    frontier.push(Node(initial, None))
    # explored is where we've been
    explored: Set[T] = {initial}

    # keep going while there is more to explore
    while not frontier.empty:
        current_node: Node[T] = frontier.pop()
        current_state: T = current_node.state
        # if we found the goal, we're done
        if goal_test(current_state):
            return current_node
        
        # 深度优先的策略：
        # 添加当前所有的子节点，但是只处理其中一个子节点
        # 如果该子节点有路可走，会继续下去到孙节点
        # 如果该子节点无路可走，则会返回，继续处理下一个子节点
        # check where we can go next and haven't explored
        for child in successors(current_state):
            #  explored 记录了前面节点可以到达的节点，如果该子节点可以到达，理论上从前面节点直接过去，所以此处不再处理
            if child in explored:  # skip children we already explored
                continue
            # 所有子节点都要添加到 explored 集合，以避免后续的节点再返回
            explored.add(child)
            # 所有的子节点都添加到 frontier 栈中，以便后续处理
            frontier.push(Node(child, current_node))

    return None  # went through everything and never found goal

# 利用 parent 反推出正确路径
def node_to_path(node:Node[T]) -> List[T]: 
    path: List[T] = [node.state]
    # work backwards from end to front
    while node.parent is not None:
        node = node.parent
        path.append(node.state)
    path.reverse()
    return path

# 队列，先进先出，使用deque实现
class Queue(Generic[T]):
    def __init__(self) -> None:
         self._container: Deque[T] = Deque()

    @property
    def empty(self) -> bool:
        return not self._container# not is true for empty container
    def push(self, item: T) -> None:
            self._container.append(item)

    def pop(self) -> T:
        return self._container.popleft()# FIFO
    def __repr__(self) -> str:
        return repr(self._container)
    
# 广度优先搜索
def bfs(initial: T, goal_test: Callable[[T], bool], successors: Callable[[T], List[T]]) -> Optional[Node[T]]:
    # frontier is where we've yet to go
    frontier: Queue[Node[T]] = Queue()
    frontier.push(Node(initial, None))
    # explored is where we've been
    explored: Set[T] = {initial}

    # keep going while there is more to explore
    while not frontier.empty:
        current_node: Node[T] = frontier.pop()
        current_state: T = current_node.state

        # if we found the goal, we're done
        if goal_test(current_state):
            return current_node
        # check where we can go next and haven't explored
        for child in successors(current_state):
            if child in explored:# skip children we already explored
                continue
            explored.add(child)
            frontier.push(Node(child, current_node))
    return None  # went through everything and never found goal


class PriorityQueue(Generic[T]):
    def __init__(self) -> None:
        self._container: List[T] = []

    @property
    def empty(self) -> bool:
        return not self._container# not is true for empty container
    def push(self, item: T) -> None:
        heappush(self._container, item)# in by priority
    def pop(self) -> T:
        return heappop(self._container)# out by priority
    def __repr__(self) -> str:
        return repr(self._container)
    
def astar(initial: T, goal_test: Callable[[T], bool], successors: Callable[[T], List[T]], heuristic: Callable[[T], float]) -> Optional[Node[T]]:
    # frontier is where we've yet to go
    frontier: PriorityQueue[Node[T]] = PriorityQueue()
    frontier.push(Node(initial, None, 0.0, heuristic(initial)))
    # explored is where we've been
    explored: Dict[T, float] = {initial: 0.0}

    # keep going while there is more to explore
    while not frontier.empty:
        current_node: Node[T] = frontier.pop()
        current_state: T = current_node.state
        # if we found the goal, we're done
        if goal_test(current_state):
            return current_node
        # check where we can go next and haven't explored
        for child in successors(current_state):
            # 1 assumes a grid, need a cost function for more sophisticated apps
            new_cost: float = current_node.cost + 1  
            if child not in explored or explored[child] > new_cost:
                explored[child] = new_cost
                frontier.push(Node(child, current_node, new_cost, heuristic(child)))
    return None# went through everything and never found goal


